import asyncio
import io
from fastmcp import FastMCP
from fastmcp import Client
import pandas as pd
from typing import Literal
import yfinance as yf
import talib

# 使用到yfinance接口，测试时获取不到数据

# 需要安装依赖：
# pip install talib

mcp = FastMCP("demo.mcp")

# 计算指定股票代码的技术指标（SMA, EMA, RSI, MACD, BBANDS）
@mcp.tool()
def calculate_technical_indicator(
    ticker: str,
    indicator: Literal["SMA", "EMA", "RSI", "MACD", "BBANDS"],
    period: Literal["1mo", "3mo", "6mo", "1y", "2y", "5y"] = "1y",
    timeperiod: int = 14,
    fastperiod: int = 12,
    slowperiod: int = 26,
    signalperiod: int = 9,
    nbdev: int = 2,
    matype: int = 0,
    num_results: int = 100
) -> dict[str, str]:
    """
    Calculate technical indicators for a stock and return CSV strings.

    Returns:
        - price_data: CSV of historical OHLCV
        - indicator_data: CSV of calculated indicator values
    """
    # Validate ticker
    if not ticker or not isinstance(ticker, str):
        raise ValueError("Ticker must be a non-empty string")
    ticker = ticker.upper().strip()

    # Fetch historical prices (daily)
    stock = yf.Ticker(ticker)
    history = stock.history(period=period, interval="1d")
    if history is None or history.empty or 'Close' not in history.columns:
        raise ValueError(f"No valid historical data found for {ticker}")

    close_prices = history['Close'].values

    # Ensure sufficient data
    min_required = {
        "SMA": timeperiod, "EMA": timeperiod*2, "RSI": timeperiod+1,
        "MACD": slowperiod + signalperiod, "BBANDS": timeperiod
    }.get(indicator, timeperiod)
    if len(close_prices) < min_required:
        raise ValueError(f"Insufficient data for {indicator}")

    # Calculate indicator
    if indicator == "SMA":
        values = talib.SMA(close_prices, timeperiod=timeperiod)
        indicator_values = {"sma": values}
    elif indicator == "EMA":
        values = talib.EMA(close_prices, timeperiod=timeperiod)
        indicator_values = {"ema": values}
    elif indicator == "RSI":
        values = talib.RSI(close_prices, timeperiod=timeperiod)
        indicator_values = {"rsi": values}
    elif indicator == "MACD":
        macd, signal, hist = talib.MACD(close_prices, fastperiod=fastperiod,
                                        slowperiod=slowperiod, signalperiod=signalperiod)
        indicator_values = {"macd": macd, "signal": signal, "histogram": hist}
    elif indicator == "BBANDS":
        upper, middle, lower = talib.BBANDS(close_prices, timeperiod=timeperiod,
                                            nbdevup=nbdev, nbdevdn=nbdev, matype=matype)
        indicator_values = {"upper_band": upper, "middle_band": middle, "lower_band": lower}

    # Limit to num_results
    if num_results > 0:
        history = history.tail(num_results)
    history_df = history.reset_index()
    history_df['Date'] = pd.to_datetime(history_df['Date']).dt.strftime('%Y-%m-%d')

    # Build indicator DataFrame
    indicator_rows = []
    for i, date in enumerate(history_df['Date']):
        row = {"Date": date}
        for name, vals in indicator_values.items():
            vals_slice = vals[-num_results:] if num_results > 0 else vals
            row[name] = f"{vals_slice[i]:.4f}" if i < len(vals_slice) and not np.isnan(vals_slice[i]) else "N/A"
        indicator_rows.append(row)
    indicator_df = pd.DataFrame(indicator_rows)

    # Convert to CSV
    price_csv = io.StringIO()
    history_df.to_csv(price_csv, index=False)
    indicator_csv = io.StringIO()
    indicator_df.to_csv(indicator_csv, index=False)

    return {"price_data": price_csv.getvalue(), "indicator_data": indicator_csv.getvalue()}

CALCULATE_TECHNICAL_INDICATOR_SCHEMA = {
    "type": "function",
    "function": {
        "name": "calculate_technical_indicator",
        "description": "Calculate a technical indicator (SMA, EMA, RSI, MACD, BBANDS) for a given stock ticker and return CSV strings for price data and indicator values.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'NFLX')."
                },
                "indicator": {
                    "type": "string",
                    "enum": ["SMA", "EMA", "RSI", "MACD", "BBANDS"],
                    "description": "The technical indicator to calculate."
                },
                "period": {
                    "type": "string",
                    "enum": ["1mo", "3mo", "6mo", "1y", "2y", "5y"],
                    "description": "The historical data period to use. Default is '1y'."
                },
                "timeperiod": {
                    "type": "integer",
                    "description": "The time period used in the indicator calculation. Default is 14."
                },
                "fastperiod": {
                    "type": "integer",
                    "description": "Fast period for MACD calculation. Default is 12."
                },
                "slowperiod": {
                    "type": "integer",
                    "description": "Slow period for MACD calculation. Default is 26."
                },
                "signalperiod": {
                    "type": "integer",
                    "description": "Signal period for MACD calculation. Default is 9."
                },
                "nbdev": {
                    "type": "integer",
                    "description": "Number of standard deviations for Bollinger Bands. Default is 2."
                },
                "matype": {
                    "type": "integer",
                    "description": "Moving average type for Bollinger Bands. Default is 0 (SMA)."
                },
                "num_results": {
                    "type": "integer",
                    "description": "Maximum number of rows to include in the returned data. Default is 100."
                }
            },
            "required": ["ticker", "indicator"]
        }
    }
}

async def main():
    client = Client(mcp)
    async with client:
        tools = await client.list_tools()
        print("Available tools:", tools)

        result = await client.call_tool(
            "test_calculate_technical_indicator", {
                "ticker": "AAPL", 
                "indicator": "SMA", 
                "num_results": 10
         })
        print("Technical Indicator Result:", result)

if __name__ == "__main__":
    asyncio.run(main())